#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os, sys
from pathlib import Path
from datetime import datetime
from argparse import ArgumentParser, Namespace
from typing import Callable
import torch.nn as nn


class HyperParams:

    # initila letter must be unique ↓↓↓
    SHORT_HANDS = [
        'source_path',
        'model_path',
        #'resolution'
        'images',
        'white_background',
        'rasterizer',
    ]

    def __init__(self):
        ''' Data '''
        self.source_path = ''
        self.model_path = ''
        self.images = 'images'
        self.limit = 99999          # limit sample count
        self.eval = False           # split train/test
        self.resolution = -1

        ''' Model '''        
        self.n_freqs = 2
        self.feat_dim = 32
        self.add_view = True
        self.appearance_dim = 0
        
        self.lambda_dssim = 0.2
        self.lambda_low = 0.01
        self.lambda_high = 0.01
        self.densification_interval = 100
        self.opacity_reset_interval = 3000
        self.densify_from_iter = 500
        self.densify_until_iter = 15_000

        # ============= Freq split methods ==============
        self.freq_split = True
        self.split_method = 'linear_split'  # ['exp_split', 'linear_split']
        
        # linear_split
        self.d_factor = 0.1
        self.D_factor = 0.3
        self.clip_min = 0.0
        
        # exp_split
        self.scale = 0.2

        ''' Optimizer '''
        self.iterations = 30_000
        self.joint_train_from_iter = 20_000

        ''' Checkpoint '''
        self.model_path = ''
        self.test_iterations = [7_000, 30_000]
        self.save_iterations = [7_000, 30_000]
        self.checkpoint_iterations = [30_000]
        self.load = ''          # resume from checkpoint *.pth

        ''' Pipeline '''
        self.rasterizer = 'original'
        self.white_background = False
        self.random_background = True
        self.compute_cov3D_python = False
        self.debug = False      # debug rasterizer

    def send_to(self, parser:ArgumentParser):
        for key, val in vars(self).items():
            if isinstance(val, Callable): 
                continue
            args = (f'--{key}',) + ((f'-{key[0]}',) if key in self.SHORT_HANDS else tuple())
            kwargs = {'default': val}
            tval = type(val)
            if tval == bool:
                kwargs.update({'action': 'store_true'})
            elif tval in [list, tuple]:
                kwargs.update({'nargs': '+', 'type': type(val[0])})
            else:
                kwargs.update({'type': tval})
            parser.add_argument(*args, **kwargs)

    def extract_from(self, args:Namespace):
        # allow override by cmd_args
        for key, val in vars(args).items():
            if isinstance(val, Callable): 
                continue
            setattr(self, key, val)
        # postfix
        self.source_path = args.source_path
        self.model_path = args.model_path
        self.checkpoint_iterations.append(self.iterations)
        if not self.model_path:
            exp_name = self.source_path.split('/')[-1]
            self.model_path = os.path.join('output', exp_name + f"-{self.task}")
            
class OptimizationParams():
    def __init__(self, 
                 postion_lr=0.00016, 
                 feature_lr=0.0075,
                 opacity_lr=0.05,
                 scaling_lr=0.005,
                 rotation_lr=0.001,
                 color_mlp_lr=0.008,
                 appearance_lr=0.05,
                 color_activation=nn.Sigmoid(),
                 densify_grad_threshold=0.0002,
                 percent_dense=0.01,
                 min_opacity=0.01, 
                 max_screen_size=None,
                 sample_ratio=1,
                 init_scaling=1.0,
                 color_scale=1.0):
        
        self.position_lr_init = postion_lr
        self.position_lr_final = postion_lr / 100
        self.position_lr_delay_mult = 0.01
        self.position_lr_max_steps = 30_000
        
        self.feature_lr = feature_lr
        self.opacity_lr = opacity_lr
        self.scaling_lr = scaling_lr
        self.rotation_lr = rotation_lr
        
        self.percent_dense = percent_dense
        self.densify_grad_threshold = densify_grad_threshold
        self.min_opacity = min_opacity
        self.max_screen_size = max_screen_size
        self.sample_ratio = sample_ratio
        self.init_scaling = init_scaling
        self.color_activation = color_activation
        self.color_scale = color_scale
        
        self.color_mlp_lr_init = color_mlp_lr
        self.color_mlp_lr_final = color_mlp_lr / 100
        self.color_mlp_lr_delay_mult = 0.01
        self.color_mlp_lr_max_steps = 30_000
        
        self.appearance_lr_init = appearance_lr
        self.appearance_lr_final = appearance_lr / 100
        self.appearance_lr_delay_mult = 0.01
        self.appearance_lr_max_steps = 30_000
    def keyParam(self):
        return {
            "percent_dense": self.percent_dense,
            "densify_grad_threshold": self.densify_grad_threshold,
            "min_opacity": self.min_opacity,
            "max_screen_size": self.max_screen_size,
            "init_scaling": self.init_scaling,
            "sample_ratio":self.sample_ratio,
        }
low_freq_opt = OptimizationParams(postion_lr=0.00016, 
                                feature_lr=0.0075,
                                opacity_lr=0.05,
                                scaling_lr=0.005,
                                rotation_lr=0.001,
                                color_mlp_lr=0.008,
                                appearance_lr=0.05,
                                percent_dense=0.02,
                                color_activation=nn.Sigmoid(),
                                densify_grad_threshold=0.0002,
                                min_opacity=0.005,
                                max_screen_size=20,
                                sample_ratio=1,
                                init_scaling=1,
                                color_scale=1)

high_freq_opt = OptimizationParams(postion_lr=0.00016, 
                                feature_lr=0.0075,
                                opacity_lr=0.05,
                                scaling_lr=0.005,
                                rotation_lr=0.001,
                                color_mlp_lr=0.008,
                                appearance_lr=0.05,
                                percent_dense=0.005,
                                color_activation=nn.Tanh(),
                                densify_grad_threshold=0.0004,
                                min_opacity=0.01,
                                max_screen_size=20,
                                sample_ratio=1,
                                init_scaling=1,
                                color_scale=0.8)

# priority: cmd_args > cfg_args > default_hp
def get_combined_args(cmd_args:ArgumentParser, default_hp:HyperParams) -> Namespace:
    cfg_fp = os.path.join(cmd_args.model_path, 'cfg_args')
    try:
        with open(cfg_fp) as fh:
            cfgfile_string = fh.read()
        print(f'Load config from: {cfg_fp}')
    except:
        cfgfile_string = 'Namespace()'
        print(f'[Warn] missing config file at {cfg_fp}')
    cfg_args = eval(cfgfile_string)
    
    merged_dict = vars(default_hp).copy()
    # cfg_args > default_hp
    for k, v in vars(cfg_args).items():
        if v is not None:
            merged_dict[k] = v
    # cmd_args > cfg_args
    for k, v in vars(cmd_args).items():
        if v != getattr(default_hp, k):
            merged_dict[k] = v
    return Namespace(**merged_dict)

